{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Instance Segmentation with GeoAI\n",
    "\n",
    "[![image](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/opengeos/geoai/blob/main/docs/examples/instance_segmentation.ipynb)\n",
    "\n",
    "This notebook demonstrates how to use the new instance segmentation functionality in GeoAI for training models and running inference on geospatial data.\n",
    "\n",
    "## Overview\n",
    "\n",
    "Instance segmentation combines object detection and semantic segmentation to identify and segment individual objects in images. This is particularly useful for:\n",
    "\n",
    "- Building detection and segmentation\n",
    "- Vehicle counting and tracking\n",
    "- Infrastructure mapping\n",
    "- Object delineation in satellite imagery\n",
    "\n",
    "## New Functions\n",
    "\n",
    "GeoAI now provides clear wrapper functions for instance segmentation:\n",
    "\n",
    "### Training\n",
    "- `train_instance_segmentation_model()` - Train a Mask R-CNN model\n",
    "\n",
    "### Inference\n",
    "- `instance_segmentation()` - Run inference on a single GeoTIFF\n",
    "- `instance_segmentation_batch()` - Run inference on multiple GeoTIFFs\n",
    "\n",
    "### Model Creation\n",
    "- `get_instance_segmentation_model()` - Create a Mask R-CNN model with custom parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install packages\n",
    "To use the new functionality, ensure the required packages are installed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %pip install geoai-py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import geoai\n",
    "import os\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "First, let's check our environment and set up paths."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check if CUDA is available\n",
    "device = geoai.get_device()\n",
    "print(f\"Using device: {device}\")\n",
    "\n",
    "# Set up paths\n",
    "out_folder = \"instance_segmentation_buildings\"\n",
    "models_dir = Path(out_folder) / \"models\"\n",
    "output_dir = Path(out_folder) / \"output\"\n",
    "\n",
    "# Create directories if they don't exist\n",
    "models_dir.mkdir(parents=True, exist_ok=True)\n",
    "output_dir.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "print(f\"Working directory: {out_folder}\")\n",
    "print(f\"Models will be saved to: {models_dir}\")\n",
    "print(f\"Output will be saved to: {output_dir}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download sample data\n",
    "\n",
    "We'll use the same dataset as the semantic segmentation example for consistency."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_raster_url = (\n",
    "    \"https://huggingface.co/datasets/giswqs/geospatial/resolve/main/naip_rgb_train.tif\"\n",
    ")\n",
    "train_vector_url = \"https://huggingface.co/datasets/giswqs/geospatial/resolve/main/naip_train_buildings.geojson\"\n",
    "test_raster_url = (\n",
    "    \"https://huggingface.co/datasets/giswqs/geospatial/resolve/main/naip_test.tif\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_raster_path = geoai.download_file(train_raster_url)\n",
    "train_vector_path = geoai.download_file(train_vector_url)\n",
    "test_raster_path = geoai.download_file(test_raster_url)\n",
    "\n",
    "print(f\"Downloaded training raster: {train_raster_path}\")\n",
    "print(f\"Downloaded training vector: {train_vector_path}\")\n",
    "print(f\"Downloaded test raster: {test_raster_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize sample data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geoai.get_raster_info(train_raster_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geoai.view_vector_interactive(train_vector_path, tiles=train_raster_url)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geoai.view_raster(test_raster_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create training data\n",
    "\n",
    "We'll create training tiles for instance segmentation. Note that for instance segmentation, we need to ensure each building instance has a unique pixel value in the label."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create training tiles\n",
    "tiles = geoai.export_geotiff_tiles(\n",
    "    in_raster=train_raster_path,\n",
    "    out_folder=out_folder,\n",
    "    in_class_data=train_vector_path,\n",
    "    tile_size=512,\n",
    "    stride=256,\n",
    "    buffer_radius=0,\n",
    ")\n",
    "\n",
    "print(f\"Created {len(tiles)} training tiles\")\n",
    "print(f\"Images saved to: {out_folder}/images\")\n",
    "print(f\"Labels saved to: {out_folder}/labels\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Model Creation\n",
    "\n",
    "Let's create an instance segmentation model with custom parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a model for binary segmentation (background + buildings)\n",
    "model = geoai.get_instance_segmentation_model(\n",
    "    num_classes=2,  # background + buildings\n",
    "    num_channels=3,  # RGB channels\n",
    "    pretrained=True,\n",
    ")\n",
    "\n",
    "print(f\"Model created with {sum(p.numel() for p in model.parameters())} parameters\")\n",
    "print(f\"Model device: {next(model.parameters()).device}\")\n",
    "print(f\"Model type: {type(model)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Training Instance Segmentation Model\n",
    "\n",
    "Now let's train the instance segmentation model using our prepared data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training configuration\n",
    "training_config = {\n",
    "    \"images_dir\": f\"{out_folder}/images\",\n",
    "    \"labels_dir\": f\"{out_folder}/labels\",\n",
    "    \"output_dir\": str(models_dir),\n",
    "    \"num_classes\": 2,  # background + buildings\n",
    "    \"num_channels\": 3,  # RGB\n",
    "    \"batch_size\": 2,  # Small batch size for demo\n",
    "    \"num_epochs\": 20,  # Few epochs for demo\n",
    "    \"learning_rate\": 0.005,\n",
    "    \"val_split\": 0.2,\n",
    "    \"visualize\": True,\n",
    "    \"device\": device,\n",
    "    \"verbose\": True,\n",
    "}\n",
    "\n",
    "print(\"Training configuration:\")\n",
    "for key, value in training_config.items():\n",
    "    print(f\"  {key}: {value}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train the model\n",
    "print(\"Starting training...\")\n",
    "geoai.train_instance_segmentation_model(**training_config)\n",
    "print(\"Training completed!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geoai.plot_performance_metrics(\n",
    "    history_path=str(models_dir / \"training_history.pth\"),\n",
    "    figsize=(15, 5),\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Running Inference\n",
    "\n",
    "Once we have a trained model, we can run inference on new images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define paths\n",
    "model_path = str(models_dir / \"best_model.pth\")\n",
    "output_path = str(output_dir / \"instance_segmentation_result.tif\")\n",
    "\n",
    "# Check if model exists\n",
    "if os.path.exists(model_path):\n",
    "    print(f\"Model found at: {model_path}\")\n",
    "else:\n",
    "    print(f\"Model not found at: {model_path}\")\n",
    "    print(\"Please ensure training completed successfully\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Single image inference with improved parameters\n",
    "inference_config = {\n",
    "    \"input_path\": test_raster_path,\n",
    "    \"output_path\": output_path,\n",
    "    \"model_path\": model_path,\n",
    "    \"window_size\": 512,\n",
    "    \"overlap\": 128,  # Reduced overlap to minimize artifacts\n",
    "    \"confidence_threshold\": 0.5,\n",
    "    \"batch_size\": 2,\n",
    "    \"num_channels\": 3,\n",
    "    \"num_classes\": 2,\n",
    "    \"device\": device,\n",
    "}\n",
    "\n",
    "print(\"Running inference with sliding window processing...\")\n",
    "result_path, inference_time = geoai.instance_segmentation(**inference_config)\n",
    "print(f\"Inference completed in {inference_time:.2f} seconds!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geoai.view_raster(\n",
    "    output_path, nodata=0, colormap=\"tab20\", opacity=0.7, basemap=test_raster_url\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](https://github.com/user-attachments/assets/e2555842-3060-4f5a-8481-09e77f05761a)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Vectorize and Visualize Results\n",
    "\n",
    "Convert the predicted mask to vector format for better visualization and analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_vector_path = \"building_predictions.geojson\"\n",
    "gdf = geoai.orthogonalize(output_path, output_vector_path, epsilon=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add geometric properties\n",
    "gdf_props = geoai.add_geometric_properties(gdf, area_unit=\"m2\", length_unit=\"m\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Interactive visualization with area information\n",
    "geoai.view_vector_interactive(gdf_props, column=\"area_m2\", tiles=test_raster_url)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Filter out small buildings and visualize\n",
    "gdf_filtered = gdf_props[(gdf_props[\"area_m2\"] > 50)]\n",
    "print(f\"Buildings after filtering (area > 50 m²): {len(gdf_filtered)}\")\n",
    "\n",
    "geoai.view_vector_interactive(gdf_filtered, column=\"area_m2\", tiles=test_raster_url)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a split map comparison\n",
    "geoai.create_split_map(\n",
    "    left_layer=gdf_filtered,\n",
    "    right_layer=test_raster_url,\n",
    "    left_args={\"style\": {\"color\": \"red\", \"fillOpacity\": 0.3}},\n",
    "    basemap=test_raster_url,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Key Parameters Guide\n",
    "\n",
    "Here are the key parameters you can adjust:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This notebook demonstrated the new instance segmentation functionality in GeoAI:\n",
    "\n",
    "1. **Model Creation**: Created Mask R-CNN models with custom parameters\n",
    "2. **Training**: Trained an instance segmentation model on building data\n",
    "3. **Inference**: Ran inference on test images\n",
    "4. **Visualization**: Converted results to vectors and visualized them\n",
    "5. **Analysis**: Compared with semantic segmentation approaches\n",
    "\n",
    "The new functions provide a cleaner API while maintaining backward compatibility with existing code. They're built on top of the robust MaskRCNN implementation already present in GeoAI.\n",
    "\n",
    "### Available Functions:\n",
    "- `geoai.train_instance_segmentation_model()` - Train Mask R-CNN models\n",
    "- `geoai.instance_segmentation()` - Single image inference\n",
    "- `geoai.instance_segmentation_batch()` - Batch processing\n",
    "- `geoai.get_instance_segmentation_model()` - Create custom models"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "geo",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
